package com.aos.iot.socket.common.scheduler;

import io.netty.channel.ChannelHandlerContext;
import io.netty.util.HashedWheelTimer;
import io.netty.util.Timeout;
import io.netty.util.internal.PlatformDependent;

import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;

public class HashedWheelTimeoutScheduler implements CancelableScheduler {

    private final ConcurrentMap<String, Timeout> scheduledFutures = PlatformDependent.newConcurrentHashMap();

    private final HashedWheelTimer executorService;

    private volatile ChannelHandlerContext ctx;

    public HashedWheelTimeoutScheduler() {
        executorService = new HashedWheelTimer();
    }

    public HashedWheelTimeoutScheduler(ThreadFactory threadFactory) {
        executorService = new HashedWheelTimer(threadFactory);
    }

    @Override
    public void update(ChannelHandlerContext ctx) {
        this.ctx = ctx;
    }

    @Override
    public void cancel(String key) {
        Timeout timeout = scheduledFutures.remove(key);
        if (timeout != null) {
            timeout.cancel();
        }
    }

    @Override
    public void scheduleCallback(String key, Runnable runnable, long delay, TimeUnit unit) {
        Timeout timeout = executorService.newTimeout(timeout1 -> ctx.executor().execute(() -> {
            try {
                runnable.run();
            } finally {
                scheduledFutures.remove(key);
            }
        }), delay, unit);
        replaceScheduledFuture(key, timeout);
    }

    @Override
    public void schedule(Runnable runnable, long delay, TimeUnit unit) {
        executorService.newTimeout(timeout -> runnable.run(), delay, unit);
    }

    @Override
    public void schedule(String key, Runnable runnable, long delay, TimeUnit unit) {
        Timeout timeout = executorService.newTimeout(timeout1 -> {
            try {
                runnable.run();
            } finally {
                scheduledFutures.remove(key);
            }
        }, delay, unit);
        replaceScheduledFuture(key, timeout);
    }

    @Override
    public void shutdown() {
        executorService.stop();
    }

    private void replaceScheduledFuture(final String key, final Timeout newTimeout) {
        final Timeout oldTimeout;
        if (newTimeout.isExpired()) {
            oldTimeout = scheduledFutures.remove(key);
        } else {
            oldTimeout = scheduledFutures.put(key, newTimeout);
        }
        if (oldTimeout != null) {
            oldTimeout.cancel();
        }
    }
}
